#include "codec/ffmpeg_encoder.h"
extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libswscale/swscale.h>
}
#include <stdexcept>

struct FFmpegEncoder::Impl {
    AVFormatContext* fmt_ctx = nullptr;
    AVCodecContext* codec_ctx = nullptr;
    AVStream* stream = nullptr;
    AVFrame* frame = nullptr;
    AVPacket* pkt = nullptr;
    SwsContext* sws_ctx = nullptr;
    int width = 0, height = 0;
    int fps = 25;
    bool opened = false;
};

FFmpegEncoder::FFmpegEncoder() : impl_(new Impl) {
    av_register_all();
    avformat_network_init();
}

FFmpegEncoder::~FFmpegEncoder() {
    close();
    delete impl_;
}

bool FFmpegEncoder::open(const std::string& output_url, int width, int height, int fps, int bitrate) {
    close();
    auto& e = *impl_;
    e.width = width;
    e.height = height;
    e.fps = fps;
    avformat_alloc_output_context2(&e.fmt_ctx, nullptr, nullptr, output_url.c_str());
    if (!e.fmt_ctx) return false;
    AVCodec* codec = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
    if (!codec) return false;
    e.stream = avformat_new_stream(e.fmt_ctx, codec);
    e.codec_ctx = avcodec_alloc_context3(codec);
    e.codec_ctx->codec_id = AV_CODEC_ID_MJPEG;
    e.codec_ctx->width = width;
    e.codec_ctx->height = height;
    e.codec_ctx->pix_fmt = AV_PIX_FMT_YUVJ420P;
    e.codec_ctx->time_base = {1, fps};
    e.codec_ctx->bit_rate = bitrate;
    if (e.fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
        e.codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
    if (avcodec_open2(e.codec_ctx, codec, nullptr) < 0) return false;
    avcodec_parameters_from_context(e.stream->codecpar, e.codec_ctx);
    if (!(e.fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
        if (avio_open(&e.fmt_ctx->pb, output_url.c_str(), AVIO_FLAG_WRITE) < 0) return false;
    }
    avformat_write_header(e.fmt_ctx, nullptr);
    e.frame = av_frame_alloc();
    e.frame->format = e.codec_ctx->pix_fmt;
    e.frame->width = width;
    e.frame->height = height;
    av_frame_get_buffer(e.frame, 0);
    e.pkt = av_packet_alloc();
    e.sws_ctx = sws_getContext(width, height, AV_PIX_FMT_BGR24, width, height, e.codec_ctx->pix_fmt, SWS_BILINEAR, nullptr, nullptr, nullptr);
    e.opened = true;
    return true;
}

bool FFmpegEncoder::writeFrame(const cv::Mat& frame) {
    auto& e = *impl_;
    if (!e.opened) return false;
    const int stride[] = {static_cast<int>(frame.step[0])};
    sws_scale(e.sws_ctx, &frame.data, stride, 0, e.height, e.frame->data, e.frame->linesize);
    e.frame->pts++;
    int ret = avcodec_send_frame(e.codec_ctx, e.frame);
    if (ret < 0) return false;
    while (ret >= 0) {
        ret = avcodec_receive_packet(e.codec_ctx, e.pkt);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) break;
        if (ret < 0) return false;
        av_packet_rescale_ts(e.pkt, e.codec_ctx->time_base, e.stream->time_base);
        e.pkt->stream_index = e.stream->index;
        av_interleaved_write_frame(e.fmt_ctx, e.pkt);
        av_packet_unref(e.pkt);
    }
    return true;
}

void FFmpegEncoder::close() {
    auto& e = *impl_;
    if (e.opened && e.fmt_ctx) av_write_trailer(e.fmt_ctx);
    if (e.frame) av_frame_free(&e.frame);
    if (e.pkt) av_packet_free(&e.pkt);
    if (e.codec_ctx) avcodec_free_context(&e.codec_ctx);
    if (e.fmt_ctx) {
        if (!(e.fmt_ctx->oformat->flags & AVFMT_NOFILE) && e.fmt_ctx->pb)
            avio_closep(&e.fmt_ctx->pb);
        avformat_free_context(e.fmt_ctx);
    }
    if (e.sws_ctx) sws_freeContext(e.sws_ctx);
    e.opened = false;
}

bool FFmpegEncoder::isOpened() const { return impl_->opened; }
